`include "all_define.vh"

module exe_stage (
  input clk,
  input rst_n,

  input ctrl_exe_flush,
  input ctrl_exe_stall,

  input                             id_exe_vld,
  input                             id_exe_is_lui,
  input                             id_exe_is_bench,
  input                             id_exe_is_jal,
  input                             id_exe_is_jalr,
  input      [31:0]                 id_exe_bench_target,
  input      [31:0]                 id_exe_pc,
  input      [31:0]                 id_exe_pc_inc,
  input                             id_exe_op1_is_pc,
  input                             id_exe_op2_is_imm,
  input                             id_exe_is_unsigned,
  input      [`ALU_TYPE_WIDTH-1:0]  id_exe_alu_type,
  input      [`ALU_OPSEL_WIDTH-1:0] id_exe_alu_opsel,
  input      [31:0]                 id_exe_imm_data,
  input      [31:0]                 id_exe_rs1_data,
  input      [31:0]                 id_exe_rs2_data,

  input                             id_exe_mem_ena,
  input      [`MEM_DTYPE_WIDTH-1:0] id_exe_mem_dtype,
  input                             id_exe_mem_rwn,

  input                             id_exe_wb_ena,
  input      [4:0]                  id_exe_wb_rd,

  output reg                        exe_mem_vld,
  output reg                        exe_mem_opena,
  output reg                        exe_mem_rwn,
  output reg [`MEM_DTYPE_WIDTH-1:0] exe_mem_dtype,
  output reg [31:0]                 exe_mem_addr,
  output reg [31:0]                 exe_mem_datain,

  output reg                        exe_mem_wb_ena,
  output reg [4:0]                  exe_mem_wb_rd,
  output reg [31:0]                 exe_mem_wb_rd_data,

  output reg                        exe_jmp_req,
  output reg [31:0]                 exe_jmp_target,

  output                            exe_bypass_rd_vld,
  output                            exe_bypass_wait_load,
  output     [4:0]                  exe_bypass_rd,
  output     [31:0]                 exe_bypass_rd_data
);

  //////////////////////////
  //     ALU Logic        //
  /////////////////////////
  reg  [31:0] alu_out_s;

  wire        alu_shift_sel_s;
  wire [31:0] alu_in1_s, alu_in2_s;
  wire [31:0] alu_logic_in1_s, alu_logic_in2_s;
  wire [31:0] alu_cmp_in1_s, alu_cmp_in2_s;
  wire [31:0] alu_arith_in1_s, alu_arith_in2_s;
  wire [31:0] alu_shift_in1_s;
  wire [4:0]  alu_shift_in2_s;

  reg  [31:0] alu_logic_out_s;
  wire [31:0] alu_cmp_out_s;
  wire [31:0] alu_arith_out_s;
  wire [31:0] alu_shift_out_s;

  wire [31:0] alu_arith_in2_invert;
  wire        alu_arith_invert_en;

  assign alu_in1_s = id_exe_op1_is_pc?  id_exe_pc: id_exe_rs1_data;
  assign alu_in2_s = id_exe_op2_is_imm? id_exe_imm_data: id_exe_rs2_data;
  assign alu_shift_sel_s = id_exe_alu_type == `ALU_TYPE_SHIFT;

  assign alu_logic_in1_s = alu_in1_s;
  assign alu_logic_in2_s = alu_in2_s;
  assign alu_cmp_in1_s   = alu_in1_s;
  assign alu_cmp_in2_s   = alu_in2_s;
  assign alu_arith_in1_s = alu_in1_s;
  assign alu_arith_in2_s = alu_in2_s;
  assign alu_shift_in1_s = alu_shift_sel_s? alu_in1_s: 32'b0;
  assign alu_shift_in2_s = alu_shift_sel_s? alu_in2_s[4:0]: 5'b0;

  always @(*) begin
    case (id_exe_alu_opsel)
      `ALU_LOGIC_AND: alu_logic_out_s = alu_logic_in1_s & alu_logic_in2_s;
      `ALU_LOGIC_OR:  alu_logic_out_s = alu_logic_in1_s | alu_logic_in2_s;
      `ALU_LOGIC_XOR: alu_logic_out_s = alu_logic_in1_s ^ alu_logic_in2_s;
      default:        alu_logic_out_s = 32'b0;
    endcase
  end

  wire alu_cmp_eq;
  wire alu_cmp_lt;
  wire alu_cmp_opout;
  wire alu_cmp_inv_out;
  assign alu_cmp_eq = (alu_cmp_in1_s == alu_cmp_in2_s);
  assign alu_cmp_lt = id_exe_is_unsigned? (alu_cmp_in1_s < alu_cmp_in2_s):
                                          ($signed(alu_cmp_in1_s) < $signed(alu_cmp_in2_s));
  assign alu_cmp_opout   = id_exe_alu_opsel[1]? alu_cmp_lt: alu_cmp_eq;
  assign alu_cmp_inv_out = id_exe_alu_opsel[0]? ~alu_cmp_opout: alu_cmp_opout;
  assign alu_cmp_out_s = {31'b0, alu_cmp_inv_out};

  assign alu_arith_in2_invert = (~alu_arith_in2_s) + 1'b1;
  assign alu_arith_invert_en = (id_exe_alu_opsel == `ALU_ARITH_SUB);
  assign alu_arith_out_s = alu_arith_in1_s + 
                           (alu_arith_invert_en? alu_arith_in2_invert: alu_arith_in2_s);

  wire [31:0] alu_shift_ll;
  wire [31:0] alu_shift_rl;
  wire [31:0] alu_shift_ra;
  assign alu_shift_ll    = (alu_shift_in1_s << alu_shift_in2_s[4:0]);
  assign alu_shift_rl    = (alu_shift_in1_s >> alu_shift_in2_s[4:0]);
  assign alu_shift_ra    = ($signed(alu_shift_in1_s) >>> alu_shift_in2_s[4:0]);
  assign alu_shift_out_s = id_exe_alu_opsel[0]?
                           (id_exe_alu_opsel[1]? alu_shift_ra: alu_shift_rl):
                                                 alu_shift_ll;

  always @(*) begin
    case (id_exe_alu_type)
      `ALU_TYPE_LOGIC: alu_out_s = alu_logic_out_s;
      `ALU_TYPE_CMP  : alu_out_s = alu_cmp_out_s;
      `ALU_TYPE_ARITH: alu_out_s = alu_arith_out_s;
      `ALU_TYPE_SHIFT: alu_out_s = alu_shift_out_s;
    endcase
  end


  ///////////////////////
  //    Stage Output   //
  ///////////////////////

  wire update_en;
  assign update_en = ~ctrl_exe_stall;

  always @(posedge clk, negedge rst_n) begin
    if (~rst_n) begin
      exe_mem_vld <= 1'b0;
    end else if (update_en) begin
      exe_mem_vld <= ~ctrl_exe_flush & id_exe_vld;
    end
  end

  always @(posedge clk, negedge rst_n) begin
    if (~rst_n) begin
      exe_mem_opena <= 1'b0;
      exe_mem_rwn <= 1'b1;
    end else if (update_en) begin
      exe_mem_opena <= id_exe_vld & id_exe_mem_ena;
      exe_mem_rwn <= id_exe_mem_rwn;
    end
  end

  always @(posedge clk, negedge rst_n) begin
    if (~rst_n) begin
      exe_mem_dtype <= `MEM_DTYPE_WORD;
    end else if (update_en) begin
      if (id_exe_vld) begin
        exe_mem_dtype <= id_exe_mem_dtype;
      end else begin
        exe_mem_dtype <= `MEM_DTYPE_WORD;
      end
    end
  end

  always @(posedge clk, negedge rst_n) begin
    if (~rst_n) begin
      exe_mem_addr <= 32'b0;
    end else if (update_en) begin
      if (id_exe_mem_ena) begin
        exe_mem_addr <= alu_out_s;
      end
    end
  end

  always @(posedge clk, negedge rst_n) begin
    if (~rst_n) begin
      exe_mem_datain <= 32'b0;
    end else if (update_en) begin
      if (id_exe_mem_ena & ~id_exe_mem_rwn) begin
        exe_mem_datain <= id_exe_rs2_data;
      end
    end
  end

  always @(posedge clk, negedge rst_n) begin
    if (~rst_n) begin
      exe_mem_wb_ena <= 1'b0;
    end else if (update_en) begin
      exe_mem_wb_ena <= id_exe_wb_ena;
    end
  end

  always @(posedge clk, negedge rst_n) begin
    if (~rst_n) begin
      exe_mem_wb_rd <= 5'b0;
    end else if (update_en) begin
      exe_mem_wb_rd <= id_exe_wb_rd;
    end
  end

  reg [31:0] wb_data;
  always @(*) begin
    wb_data = 32'b0;
    if (id_exe_is_lui) begin
      wb_data = id_exe_imm_data;
    end else if (id_exe_is_jal | id_exe_is_jalr) begin
      wb_data = id_exe_pc_inc;
    end else if (id_exe_wb_ena) begin
      wb_data = alu_out_s;
    end
  end

  always @(posedge clk, negedge rst_n) begin
    if (~rst_n) begin
      exe_mem_wb_rd_data <= 32'b0;
    end else if (update_en) begin
      exe_mem_wb_rd_data <= wb_data;
    end
  end

  wire   jmp_req_vld;
  assign jmp_req_vld = ~(ctrl_exe_flush | ctrl_exe_stall);
  always @(posedge clk, negedge rst_n) begin
    if (~rst_n) begin
      exe_jmp_req <= 1'b0;
    end else if (update_en) begin
      exe_jmp_req <= 1'b0;
      if (jmp_req_vld) begin
        if (id_exe_is_jalr) begin
          exe_jmp_req <= 1'b1;
        end else if (id_exe_is_bench) begin
          exe_jmp_req <= alu_out_s[0];
        end
      end
    end
  end

  always @(posedge clk, negedge rst_n) begin
    if (~rst_n) begin
      exe_jmp_target <= 32'b0;
    end else if (update_en) begin
      exe_jmp_target <= 32'b0;
      if (id_exe_is_jalr) begin
        exe_jmp_target <= alu_out_s;
      end else if (id_exe_is_bench & alu_out_s[0]) begin
        exe_jmp_target <= id_exe_bench_target;
      end
    end
  end

  assign exe_bypass_rd_vld = id_exe_vld & id_exe_wb_ena;
  assign exe_bypass_wait_load = id_exe_mem_ena & id_exe_mem_rwn;
  assign exe_bypass_rd     = id_exe_wb_rd;
  assign exe_bypass_rd_data= wb_data;

endmodule

